import numpy as np
import pandas as pd
import h5py
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
import os
import Cooking
import random


RAW_DATA_DIR = 'data_raw/'
COOKED_DATA_DIR = 'data_cooked/'
DATA_FOLDERS = ['normal_1','normal_2','normal_3','normal_4','normal_5','normal_6',
                'swerve_1','swerve_2','swerve_3']
FIGURE_SIZE = (10, 10)


sample_tsv_path = os.path.join(RAW_DATA_DIR,'normal_1/airsim_rec.txt')
sample_tsv = pd.read_csv(sample_tsv_path, sep='\t')
print(sample_tsv.head())


sample_image_path = os.path.join(RAW_DATA_DIR,'normal_1/images/img_0.png')
sample_image = Image.open(sample_image_path)
plt.title('Sample Image')
plt.imshow(sample_image)
plt.show()

sample_image_roi = sample_image.copy()
#draw line
fillcolor=(255, 0, 0)
draw = ImageDraw.Draw(sample_image_roi)
points = [(1,76), (1,135), (255,135), (255,76)]
for i in range(0,len(points)):
    draw.line([points[i],points[(i+1)%len(points)]], fill=fillcolor, width=3)
del draw

plt.title('Image with sample ROI')
plt.imshow(sample_image_roi)
plt.show()

full_path_raw_folder = [os.path.join(RAW_DATA_DIR, f) for f in DATA_FOLDERS]

dataframe = []
for folder in full_path_raw_folder:
    current_dataframe = pd.read_csv(os.path.join(folder,'airsim_rec.txt'),sep='\t')
    current_dataframe['Folder'] = folder
    dataframe.append(current_dataframe)

dataset = pd.concat(dataframe, axis=0)
print('Number of data points: {0}'.format(dataset.shape))
print(dataset.head())

min_idx = 100
max_idx = 1100
normal_idx = dataset['Folder'].apply(lambda v:'normal_1' in v)
swerve_idx = dataset['Folder'].apply(lambda v:'swerve_1' in v)
steering_angles_normal_l = dataset[normal_idx]['Steering'][min_idx:max_idx]
steering_angles_swerve_l = dataset[swerve_idx]['Steering'][min_idx:max_idx]

plot_idx = [i for i in range(min_idx,max_idx)]
fig = plt.figure(figsize=FIGURE_SIZE)
ax = fig.add_subplot(111)

ax.scatter(plot_idx,steering_angles_normal_l,c='b',marker='o',label='normal_l')
ax.scatter(plot_idx,steering_angles_swerve_l,c='r',marker='o',label='swerve_l')
plt.legend(loc='upper left');
plt.title('Steering Angles for normal_l and swerve_l')
plt.xlabel('Time')
plt.ylabel('Steering Angle')
plt.show()

#pie
dataset['Is Swerve'] = dataset.apply(lambda v:'swerve' in v['Folder'],axis = 1)
grouped = dataset.groupby(by=['Is Swerve']).size().reset_index()
grouped.columns=['Is Swerve','Count']
def make_autopct(values):
    def my_autopct(percent):
        total = sum(values)
        val = int(round(percent*total/100))
        return '{0:.2f}% ({1:d})'.format(percent,val)
    return my_autopct

pie_labels = ['Normal','Swerve']
fig, ax = plt.subplots(figsize=FIGURE_SIZE)
ax.pie(grouped['Count'], labels=pie_labels, autopct=make_autopct(grouped['Count']))
plt.title('Number of data points per driving strategy')
plt.show()


#策略分布
bins = np.arange(-1, 1.05,0.05)
normal_labels = dataset[dataset['Is Swerve'] == False]['Steering']
swerve_labels = dataset[dataset['Is Swerve'] == True]['Steering']

def steering_histogram(hist_labels, title, color):
    plt.figure(figsize=FIGURE_SIZE)
    n, b, p = plt.hist(hist_labels.values, bins, facecolor=color)
    plt.xlabel('Steering Angle')
    plt.ylabel('Normalized Frequency')
    plt.title(title)
    plt.show()

steering_histogram(normal_labels, 'Normal label distribution', 'r')
steering_histogram(swerve_labels, 'Swerve label distributi on', 'g')


train_eval_test_split = [0.7, 0.2, 0.1]
full_path_raw_folder = [os.path.join(RAW_DATA_DIR, f) for f in DATA_FOLDERS]
Cooking.cook(full_path_raw_folder, COOKED_DATA_DIR, train_eval_test_split)